package org.smartboot.flow.core.component;


import org.smartboot.flow.core.Condition;
import org.smartboot.flow.core.EngineContext;
import org.smartboot.flow.core.Key;
import org.smartboot.flow.core.common.ComponentType;
import org.smartboot.flow.core.util.AssertUtil;
import org.smartboot.flow.core.visitor.ComponentVisitor;

import java.util.Map;

/**
 * @author qinluo
 * @date 2022-11-12 18:57:40
 * @since 1.0.0
 */
public class ChooseComponent<T, S> extends Component<T, S> {

    private Map<Object, Component<T, S>> branches;
    private Condition<T, S> condition;
    private Component<T, S> defaultBranch;
    private boolean allBranchWasString;

    public ChooseComponent() {
    }

    public ChooseComponent(Condition<T, S> condition, Map<Object, Component<T, S>> branches, Component<T, S> defaultBranch) {
        this.branches = branches;
        this.condition = condition;
        this.defaultBranch = defaultBranch;
    }

    public void setBranches(Map<Object, Component<T, S>> branches) {
        this.branches = branches;
    }

    public void setCondition(Condition<T, S> condition) {
        this.condition = condition;
    }

    public void setDefaultBranch(Component<T, S> defaultBranch) {
        this.defaultBranch = defaultBranch;
    }

    public void setAllBranchWasString(boolean allBranchWasString) {
        this.allBranchWasString = allBranchWasString;
    }

    @Override
    public int invoke(EngineContext<T, S> context) throws Throwable {
        context.enter(this);
        try {
            Object branch = condition.test(context);
            Component<T, S> execute = null;

            // Compatible
            if (branch != null && allBranchWasString) {
                branch = String.valueOf(branch);
            }

            if (branch != null && branches.containsKey(branch)) {
                execute = branches.get(branch);
            } else if (defaultBranch != null) {
                execute = defaultBranch;
                branch = "default";
            }

            if (execute != null) {
                context.putExt(Key.of(this), execute);
                context.enter("branch##" + branch);
                try {
                    execute.invoke(context);
                } finally {
                    context.exit("branch##" + branch);
                }
            }
        } finally {
            context.exit(this);
        }

        return 1;
    }

    @Override
    public boolean isRollbackable(EngineContext<T, S> context) {
        Component<T, S> executed = context.getExt(Key.of(this));
        return executed != null && executed.isRollbackable(context);
    }

    @Override
    public void rollback(EngineContext<T, S> context) {
        Component<T, S> executed = context.remove(Key.of(this));
        if (executed == null || !executed.isRollbackable(context)) {
            return;
        }

        context.enter(this);
        try {
            executed.rollback(context);
        } finally {
            context.exit(this);
        }

    }

    @Override
    public String describe() {
        return "choose@" + condition.describe();
    }

    @Override
    public void visit(ComponentVisitor visitor) {
        visitor.visitAttributes(attributes);
        visitor.visitCondition(condition.describe());
        visitor.visitSource(this);

        branches.forEach((k, v) -> {
            ComponentVisitor branchVisitor = visitor.visitBranch(k, v.getType(), v.getName(), v.describe());
            if (branchVisitor != null) {
                v.visit(branchVisitor);
            }
        });

        if (defaultBranch != null) {
            ComponentVisitor defaultVisitor = visitor.visitComponent(defaultBranch.getType(), defaultBranch.getName(), defaultBranch.describe());
            if (defaultVisitor != null) {
                defaultBranch.visit(defaultVisitor);
            }

        }
    }

    @Override
    public void doValidate() {
        AssertUtil.notNull(condition, "choose[" + getName() + "] condition must not be null");
        AssertUtil.notNull(branches, "choose[" + getName() + "] branch must not be null");
        AssertUtil.isTrue(branches.size() != 0, "choose[" + getName() + "] branch must not be null");
        branches.forEach((k, v) -> v.validate());

        if (defaultBranch != null) {
            defaultBranch.validate();
        }
    }

    @Override
    public ComponentType getType() {
        return ComponentType.CHOOSE;
    }
}
